# -*- coding: utf-8 -*-
import json
import os
import time
import logging
from alibabacloud_imageseg20191230.client import Client
from alibabacloud_imageseg20191230.models import SegmentHDCommonImageAdvanceRequest
from alibabacloud_tea_openapi.models import Config
from alibabacloud_tea_util.models import RuntimeOptions

from alibabacloud_viapi20230117.client import Client as asyncClient
from alibabacloud_viapi20230117.models import GetAsyncJobResultRequest
import requests

logger = logging.getLogger()


def conf(endpoint):
    config = Config(
        access_key_id=os.environ.get('ALIBABA_CLOUD_ACCESS_KEY_ID'),
        access_key_secret=os.environ.get('ALIBABA_CLOUD_ACCESS_KEY_SECRET'),
        endpoint=endpoint,
        region_id='cn-shanghai'
    )
    return config


def request_segmentation(img_path):
    """ 调用视觉智能开放平台-通用高清分割API，获取图片主体 """
    config = conf('imageseg.cn-shanghai.aliyuncs.com')
    try:
        with open(img_path, 'rb') as img:
            segment_hdcommon_image_request = SegmentHDCommonImageAdvanceRequest()
            segment_hdcommon_image_request.image_url_object = img
            runtime = RuntimeOptions()
            client = Client(config)
            response = client.segment_hdcommon_image_advance(segment_hdcommon_image_request, runtime)
            request_id = json.loads(str(response.body).replace("'", "\""))["RequestId"]
            return get_vi_result(request_id)
    except Exception as error:
        logger.error(error)
        return None


def get_vi_result(requestId):
    """ 获取图片分割结果 """
    if requestId is None:
        return None
    logger.info("RequestId: " + requestId)
    config = conf('viapi.cn-shanghai.aliyuncs.com')
    get_async_job_result_request = GetAsyncJobResultRequest(
        job_id=requestId
    )
    runtime = RuntimeOptions()
    try:
        client = asyncClient(config)
        response = client.get_async_job_result_with_options(get_async_job_result_request, runtime)
        data = json.loads(str(response.body).replace("'{", "{").replace("}'", "}").replace("'", "\""))
        logger.info(data)
        if data["Data"]["Status"] == "PROCESSING":
            time.sleep(1)
            logger.info("RequestId: " + requestId + " 异步处理中...")
            return get_vi_result(requestId)

        return data["Data"]["Result"]["imageUrl"]
    except Exception as error:
        logger.error(error)


def handle_response(response):
    error_data = response.json()
    error_code = error_data.get("code")
    error_message = error_data.get("message")
    if response.status_code != 200:
        if error_code == "InvalidURL":
            return False, "在数据检查过程中无法下载媒体资源"
        elif error_code == "InvalidParameter.DataInspection":
            return False, "在数据检查过程中无法下载媒体资源"
        elif error_code == "InvalidApiKey" or error_code == "invalid_api_key":
            return False, "百炼 API-KEY错误"
        elif error_code == "AccessDenied.Unpurchased":
            return False, "尚未未开通百炼服务，请前往模型广场开通模型服务"
        elif error_code == "Arrearage":
            return False, "当前账号已欠费，请充值后，等待10分钟后重试"
        else:
            return False, error_message
    else:
        data = response.json()
        status = data["output"]["task_status"]
        if status == "PENDING" or status == "PRE-PROCESSING" or status == "RUNNING":
            task_id = data["output"]["task_id"]
            result = get_result(task_id)
            return True, result


def generate_background(base_image_url, ref_prompt):
    """ 调用通义万相模型 """
    url = 'https://dashscope.aliyuncs.com/api/v1/services/aigc/background-generation/generation/'
    headers = {
        'X-DashScope-Async': 'enable',
        'Authorization': f'Bearer ' + os.environ.get('DASHSCOPE_API_KEY'),
        'Content-Type': 'application/json'
    }
    data = {
        "model": "wanx-background-generation-v2",
        "input": {
            "base_image_url": base_image_url,
            "ref_prompt": ref_prompt
        },
        "parameters": {
            "n": 4,
            "ref_prompt_weight": 0.5,
            "model_version": "v3"
        }
    }
    response = requests.post(url, headers=headers, json=data)
    return handle_response(response)


def get_result(task_id):
    """ 获取通义万相模型处理结果 """
    if task_id is None:
        return None
    url = 'https://dashscope.aliyuncs.com/api/v1/tasks/' + task_id
    headers = {
        'Authorization': f'Bearer ' + os.environ.get('DASHSCOPE_API_KEY'),
    }
    try:
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            data = response.json()
            task_status = data["output"]["task_status"]
            if task_status == "PENDING" or task_status == "RUNNING":
                time.sleep(1)
                logger.info("task_id: " + task_id + " 异步处理中...")
                return get_result(task_id)

            if data["output"]["task_status"] == "FAILED":
                logger.error(data["output"]["message"])
                return None
            results = data["output"]["results"]
            return [item['url'] for item in results if item['url']]
        else:
            logger.error(response.text)
            return None
    except Exception as error:
        logger.error(error)
        return None
